from typing import List

from loguru import logger
from .node import Node
import networkx as nx


def get_clean_name(name):

    if 'input' in name:
        return 'input'

    if 'conv2d' in name:
        return 'conv2d'

    if 'fused_' in name:
        name = name.replace('fused_', '')
    if '_compute_' in name:
        name = name.replace('_compute_', '')
    if 'tvmgen_default_' in name:
        name = name.replace('tvmgen_default_', '')
    if 'nn_' in name:
        name = name.replace('nn_', '')
    if name.split('_')[-1].isdigit():
        name = name[:-(len(name.split('_')[-1]) + 1)]

    name = name.split("-")[0]
    pos = len(name) - 1
    for i in range(len(name)):
        # Reversing
        pos = len(name) - i - 1
        if name[pos].isdigit() or name[pos] == '_':
            continue
        else:
            break
    name = name[:pos + 1]
    return name


def node_match(node_1, node_2) -> bool:
    return node_1['type'] == node_2['type']


class Graph:
    def __init__(self) -> None:
        self.root = Node('root', 0, None)
        self.node_set: List[Node] = []
        self.node_dict = {}
        self.nx_graph = None
        pass

    def add_node(self, node):
        if len(self.node_set) == 0:
            self.root.next_nodes = node

        self.node_set.append(node)
        self.node_dict[node.name] = node
        return self

    def find_node_by_name(self, name):
        if name in self.node_dict:
            return self.node_dict[name]
        else:
            return None

    def build_networkx_graph(self):
        networkx_graph = nx.DiGraph()
        node_num = len(self.node_set)
        networkx_graph.add_nodes_from([i for i in range(node_num)])
        for idx, node in enumerate(self.node_set):
            if node.is_extra and node.is_input:
                node.atomic_type = 'input'
            elif node.is_extra and not node.is_input:
                node.atomic_type = "p"
            networkx_graph.nodes[idx]['type'] = node.atomic_type
            networkx_graph.nodes[idx]['gt'] = get_clean_name(node.name)
            networkx_graph.nodes[idx]['name'] = node.name
            if len(node.extra_input) > 0:
                networkx_graph.nodes[idx]['max_weight_dim'] = max([len(x.out_shape) for x in node.extra_input])
            else:
                networkx_graph.nodes[idx]['max_weight_dim'] = 0

            edges = [(idx, self.node_set.index(next_node)) for next_node in node.next_nodes]
            networkx_graph.add_edges_from(edges)

        self.nx_graph = networkx_graph

    def find_sub_graph(self, sub_graph):
        # rebuild the nertwork grpah
        self.build_networkx_graph()
        dgm = nx.algorithms.isomorphism.DiGraphMatcher(self.nx_graph, sub_graph, node_match=node_match)
        res = []
        for sub_g in dgm.subgraph_isomorphisms_iter():
            res.append(sub_g)
        return res

    def merge_node(self, node_1: Node, node_2: Node, new_name: str, type_name: str) -> None:
        merged_node = Node(new_name, node_1.input_num, node_2.out_shape)
        merged_node.in_shape = node_1.in_shape
        merged_node.function_type = type_name
        merged_node.function_param = node_1.function_param
        merged_node.function_param.update(node_2.function_param)
        merged_node.extra_input = node_1.extra_input
        merged_node.extra_input.extend(node_2.extra_input)
        del self.node_dict[node_1.name]
        del self.node_dict[node_2.name]

        new_node_set = []
        for n in self.node_set:
            if n.name != node_1.name and n.name != node_2.name:
                new_node_set.append(n)
        new_node_set.append(merged_node)
        self.node_set = new_node_set

        self.node_dict[merged_node.name] = merged_node

        # Change the in-node
        for tmp_node in node_1.father_nodes:
            new_next_nodes = [merged_node] + tmp_node.next_nodes
            new_next_nodes.remove(node_1)
            tmp_node.next_nodes = new_next_nodes
            self.find_node_by_name(tmp_node.name).next_nodes = new_next_nodes
            merged_node.father_nodes.append(tmp_node)
            # tmp_node.next_nodes = new_next_nodes

        # Change the out-node
        for tmp_node in node_2.next_nodes:
            new_father_nodes = [merged_node] + tmp_node.father_nodes
            new_father_nodes.remove(node_2)
            self.find_node_by_name(tmp_node.name).father_nodes = new_father_nodes
            merged_node.add_next_node(tmp_node)

    def merge_multi_nodes(self, nodes: List[Node], new_name: str, type_name: str, attrs=None) -> None:
        node_1 = nodes[0]
        node_2 = nodes[-1]
        merged_node = Node(new_name, node_1.input_num, node_2.out_shape)
        if attrs is not None:
            merged_node.function_param = attrs

        merged_node.in_shape = node_1.in_shape
        merged_node.function_type = type_name
        name_list = [n.name for n in nodes]
        # Change the in-node
        for node in nodes:
            for tmp_node in node.father_nodes:
                if tmp_node.name in name_list:
                    continue
                new_next_nodes = [merged_node]
                for n in tmp_node.next_nodes:
                    if n.name != node.name:
                        new_next_nodes.append(n)
                tmp_node.next_nodes = new_next_nodes
                merged_node.add_father_node(tmp_node)

        # Change the out-node
        for tmp_node in node_2.next_nodes:
            new_father_nodes = [merged_node]
            for n in tmp_node.father_nodes:
                if n.name != node_2.name:
                    new_father_nodes.append(n)
            tmp_node.father_nodes = new_father_nodes
            merged_node.add_next_node(tmp_node)

        new_node_set = []
        for n in self.node_set:
            if n.name not in name_list:
                new_node_set.append(n)
        new_node_set.append(merged_node)
        self.node_set = new_node_set

        new_node_dict = {}
        for node in self.node_set:
            new_node_dict[node.name] = node
        self.node_dict = new_node_dict

    def gen_label_by_name(self):
        for node in self.node_set:
            tmp_name = node.name
            tmp_type = get_clean_name(tmp_name)
            node.atomic_type = tmp_type
            # logger.debug("Set {}'s type -> {}".format(node.name, tmp_type))

    def remove_nop(self):
        # merge_tasks = []
        nop_node_list = []
        for node in self.node_set:
            if '__nop' in node.name:
                # logger.debug("Found nop!")
                # logger.debug("Father nodes: {}".format([fn.name for fn in node.father_nodes]))
                # logger.debug("Next nodes: {}".format([nn.name for nn in node.next_nodes]))
                node.father_nodes[0].next_nodes.remove(node)
                if len(node.next_nodes) > 0:
                    node.father_nodes[0].next_nodes.append(node.next_nodes[0])
                    node.next_nodes[0].father_nodes.remove(node)
                    node.next_nodes[0].father_nodes.append(node.father_nodes[0])
                nop_node_list.append(node)
                # merge_tasks.append([node, node.next_nodes[0], node.next_nodes[0].name, node.next_nodes[0].atomic_type])
        for node in nop_node_list:
            self.node_set.remove(node)
            del self.node_dict[node.name]
